from collections import deque
import sys
import threading
import queue
import time

import numpy as np
import pyqtgraph as pg
import qtmodern.styles
import qtmodern.windows
from PySide6 import QtCore
from PySide6.QtWidgets import QApplication, QMainWindow, QGridLayout

from semg_fes_UI import Ui_MainWindow  # 导入转换得到的UI类
from data_recorder import sEMG_Data_Recorder
from FES_Controler import FES_Controler, FES_Config, FES_config_map
import multiprocessing
import data_recorder
import socket
import queue
import logging

from Sensors.SerialPortFor28BytesData_python3 import mainProcess_Task_six_axis_force

# 创建一个logger
logger = logging.getLogger("MyApp")
logger.setLevel(logging.INFO)  # 设置全局日志级别

# 创建控制台 handler 并设置级别为 INFO
console_handler = logging.StreamHandler()
console_handler.setLevel(logging.INFO)

# 设置日志格式
formatter = logging.Formatter("%(message)s")
console_handler.setFormatter(formatter)

# 添加 handler 到 logger
logger.addHandler(console_handler)

'''
全局变量
'''
data_queue = []
data_queuerms = []
command_queue = []
LocalHost = '192.168.137.1'
RemoteHost = '192.168.137.142'

target_device_address = ''

MaxplotNum = 2000  #最大绘图点数

rmsWindow_len = 8
sEMGWindow_len = 8

FES_Control_ON = False
sEMG_mode = 0  # 0:raw_sEMG, 1:rms_sEMG
Default_Channel_Intensity = 2000

debug_on = True
FES_FREQ = 200


class sensors:
    """
    其他传感器数据采集开关
    """
    six_axis_force = False
    # 添加其他传感器


# 通道电刺激状态类
class Channel_Stase:
    def __init__(self):
        self.FESChannel_State = [False for _ in range(4)]
        self.sEMGChannel_State = [True for _ in range(4)]
        # 判断FES通道是否选中了
        self.FESchanged_state = [False for _ in range(4)]

        self.channel_intensity = [Default_Channel_Intensity for _ in range(4)]
        # 生成双极性波形
        self.FES_wave = self.FES_plotwave_gen()

    def FES_plotwave_gen(self):
        """
        生成电刺激波形用于绘图
        :return: 波信号的numpy矢量
        """
        # 波形参数
        frequency = 50  # Hz，波形频率
        amplitude = 1.0  # V，波形幅度
        duration = 0.02  # s，波形持续时间
        sample_rate = 1000  # Hz，采样率
        # 生成单极性方波
        t = np.linspace(0, duration, int(sample_rate * duration), endpoint=False)
        square_wave = amplitude * np.sign(np.sin(2 * np.pi * frequency * t))

        # # 将单极性方波转换为双极性方波
        bipolar_wave = square_wave - np.mean(square_wave)
        bipolar_wave = np.append(bipolar_wave, np.zeros(30))
        # # 重复三次
        bipolar_wave = np.tile(bipolar_wave, 3)

        return bipolar_wave


channel_state = Channel_Stase()
WriteIn_Queue = multiprocessing.Queue()  # 创建用于存放数据的队列
Data_Writeflag = False

recordeprocess = None


class UDPCommunicator:
    """
    新建一个UDP传输数据类，用于使用UDP接收、处理、传输数据
    """

    def __init__(self, local_host, remote_host, port):
        self.RMS_Window_Len = 40
        self.semg_deque = [deque(maxlen=self.RMS_Window_Len) for _ in range(4)]

        self.plot_data_list = [[], [], [], []]
        self.max_plot_num = MaxplotNum

        self.sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self.server_address = (local_host, port)
        self.sock.bind(self.server_address)
        # self.sock.setblocking(False)

        self.remote_host = remote_host
        self.data_queue = queue.Queue()  # 数据-绘图队列
        self.rms_for_fes_queue = queue.Queue(maxsize=3)

        self.command_queue = queue.Queue()  # 指令队列
        self.FES_config_queue = queue.Queue()
        self.addr = None

        # --- 添加 RMS 更新控制相关的属性 ---
        self.rms_update_threshold = 2  # 每接收多少个数据块更新一次RMS
        self.data_chunk_counter = 0  # 数据块计数器
        self.last_rms_array = np.zeros(4)  # 存储上一次计算的RMS值

        self.cmd_mp = {
            # 指令哈希表
            'debug': self.switch_debug_mode,
            'testFES': self.Test_send_FESConfig,
            'ele_on': self.open_ele,
        }

    def update_rms(self, semgdata):
        """
        将numarray加入到rms_deque中并计算rms，返回rms_array (shape: (4,))
        semgdata: 输入的sEMG数据块, shape (N, 4)
        """
        rms_array = np.zeros(4)
        self.data_chunk_counter += 1  # 增加计数器
        # logger.info(f"semgdata:{semgdata}")

        for i in range(4):
            # 将新数据添加到对应通道的deque末尾
            # deque的maxlen属性会自动移除另一端的旧数据
            self.semg_deque[i].extend(semgdata[:, i])
            # logger.info(f"semgdata{i}:{semgdata[i]}")
            # 如果deque中有数据，则计算RMS
        if self.data_chunk_counter >= self.rms_update_threshold:
            for i in range(4):
                rms_array[i] = np.sqrt(np.mean(np.square(np.array(self.semg_deque[i]))))
            # 重置计数器
            self.data_chunk_counter = 0
            return rms_array.reshape(1, 4)
        else:
            return None

    def data_process_loop(self):
        while True:
            try:
                tmp_data = self.data_queue.get()
                for i in range(4):
                    """
                    先放着吧，未来再性能优化下
                    """
                    # print(i, tmp_data[:, i])
                    self.plot_data_list[i].extend(tmp_data[:, i])
                    # print("plotlen: ")
                    # print(len(self.plot_data_list[i]))
                    if len(self.plot_data_list[i]) > self.max_plot_num:
                        self.plot_data_list[i] = self.plot_data_list[i][-self.max_plot_num:]
                        # print("plotlenafter: ")
                        # print(len(self.plot_data_list[i]))
            except queue.Empty:
                pass

    def receive_data_loop_test(self):
        while True:
            # 生成从 1 到 40 的数组
            incremental_array = np.arange(1, 41)

            # 重复这个数组 4 次，形成形状为 (40, 4) 的 ndarray
            data = np.tile(incremental_array, (1, 4)).reshape(40, 4)
            print(len(data))
            print(data)
            self.data_queue.put(data)
            # 延时40ms
            time.sleep(0.04)

    def receive_and_send_loop(self):
        logger.info("等待接收消息...")
        while True:
            # 接收数据
            try:
                data, self.addr = self.sock.recvfrom(1024)
                print("收到消息:", data.decode('utf-8'), "来自:", self.addr)
                break
            except socket.error as e:
                if e.errno == socket.EAGAIN:
                    # 无数据可用，正常情况
                    continue
                    pass

        while True:
            # 接收数据
            try:
                data, self.addr = self.sock.recvfrom(1024)
                # data是以\n分隔开，将其转换为列表加入到data_queue
                # print(data)
                # 在收到外界UDO请求数据时会收到空字节数据
                if data == b'':
                    continue

                numarray = np.genfromtxt(data.decode('utf-8').splitlines(), dtype=int, delimiter=',')
                # print(numarray)
                # print(numarray.size)
                if numarray.size != 20:
                    print("hi")
                    print(numarray == 0x43)
                    print("received data")
                    continue

                global FES_Control_ON, sEMG_mode

                if Data_Writeflag == True:
                    # 将数据写入文件保存队列
                    WriteIn_Queue.put(numarray)

                rms = self.update_rms(numarray)

                if sEMG_mode == 0:  #rawsEMG模式
                    self.data_queue.put(numarray)
                elif sEMG_mode == 1:  #RMSsEMG模式
                    if rms is not None:
                        self.data_queue.put(rms)
                        logger.debug(f"收到数据: {rms}")
                if FES_Control_ON and rms is not None:
                    # logger.info(f"f{rms}数据已放入队列")
                    try:
                        self.rms_for_fes_queue.put(rms, block=False)
                    except queue.Full:
                        # 如果队列已满，则丢弃数据
                        pass
                # if rms is not None:
                #     logger.info(f"rms:{rms}")



            except socket.error as e:
                if e.errno == socket.EAGAIN:
                    # 无数据可用，正常情况
                    continue
                else:
                    print("Error receiving data:", e)
                    break

    """
    发送指令相关函数
    """

    def send_data(self, data):
        self.sock.sendto(data, self.addr)

    """
    UDP远程传输指令相关函数
    """

    def send_FES_control(self):
        """
        发送FES控制指令
        :return:
        """
        FES_is_ON = False
        if self.addr is None:
            print("未连接到目标主机")
            return
        # 发送指令
        while self.FES_config_queue.qsize() > 0:
            fes_config = self.FES_config_queue.get()
            self.send_FESConfig_to_queue(fes_config)
            logger.debug("FES is On")
            logger.debug(f"发送FES控制指令,ch:{fes_config.STI_ch}")
            FES_is_ON = True
            time.sleep(0.1)
        # 清空rms_deque
        if FES_is_ON:
            for dq in self.semg_deque:
                dq.clear()

    def Set_Waveform(self, worktime, lastingtime, holdtime):
        """
        配置波形
        :param worktime:
        :param lastingtime:
        :param holdtime:
        """
        # 发送8字节的指令数据
        print("sending:")

        hex_num = '0x43'
        hex_tail = '0x1141'
        data1 = int(hex_num, 16)
        lastingtime = int(lastingtime)
        worktime = int(worktime)
        holdtime = int(holdtime)
        print(worktime, lastingtime, holdtime)

        data_bin = (data1.to_bytes(4, 'big')
                    + worktime.to_bytes(2, 'big')
                    + lastingtime.to_bytes(1, 'big')
                    + holdtime.to_bytes(1, 'big')
                    + int(hex_tail, 16).to_bytes(8, 'big')
                    )
        self.sock.sendto(data_bin, self.addr)

    def OrderTransmit(self, data1, data2):
        # 消息内容,两个uint32_t数据
        # 将数据转换为二进制格式

        data1 = 42
        data2 = 1337
        data_bin = (data1.to_bytes(4, 'big') + data2.to_bytes(4, 'big'))
        # 向单片机发送数据
        self.sock.sendto(data_bin, self.addr)

    def CommandExecute(self, command):
        """
        执行指令:将指令字符串分解成发送指令并发送到下位机
        :param command:
        :return:
        """
        # command格式为"set ch tensity"，解析参数并调用SetTensity()函数
        command = command.strip().split()
        headfram = command[0].lower()
        if headfram == "set":
            if len(command) == 3:
                ch, tensity = int(command[1]), int(command[2])
                self.SetTensity(ch, tensity)
        elif headfram == "sw":
            if len(command) == 4:
                self.Set_Waveform(command[1], command[2], command[3])
        # elif command[0].lower() == "send":
        elif headfram == "fes":
            if len(command) == 2:
                tensity = int(command[1])
                self.send_fescongfig_intensity(tensity)  #默认电刺激0通道
            elif len(command) == 3:
                ch, tensity = int(command[1]), int(command[2])
                self.send_fescongfig_intensity(tensity, ch)
        # 其他指令
        elif headfram in self.cmd_mp:
            self.cmd_mp[headfram]()
        else:
            logger.info("Invalid command format")
            return

    def SetTensity(self, ch, tensity):
        """
        设定通道ch电刺激强度为tensity
        :param ch:-
        :param tensity:
        :return: None
        """
        # 发送8字节的指令数据
        hex_num = '0x41'
        hex_tail = '0x1141'
        data1 = int(hex_num, 16)
        tensity = min(tensity, 4095)

        data_bin = (data1.to_bytes(4, 'big') +
                    ch.to_bytes(2, 'big') + tensity.to_bytes(2, 'big') +
                    int(hex_tail, 16).to_bytes(8, 'big'))
        # 向单片机发送数据
        self.sock.sendto(data_bin, self.addr)

    def openloop_elecontrol(self):
        # global data_queuerms
        if len(self.semg_deque) == self.RMS_Window_Len:
            target = np.mean(self.semg_deque)
            if target > 1700:
                self.send_waveform(2000, 200, 50)
                print("电刺激，较大")
            elif target > 1680:
                self.send_waveform(500, 100, 25)
                print("小规模电刺激")
            elif target > 1650:
                self.send_waveform(500, 25, 25)
                print("微型电刺激")

    def send_waveform(self, worktime, lastingtime, holdtime):
        """
        直接向下位机发送电刺激波形组
        :param worktime:
        :param lastingtime:
        :param holdtime:
        :return:
        """
        hex_head = '0x43'
        hex_tail = '0x1141'
        data1 = int(hex_head, 16)
        data_bin = (data1.to_bytes(4, 'big') +
                    worktime.to_bytes(2, 'big') +
                    lastingtime.to_bytes(1, 'big') +
                    holdtime.to_bytes(1, 'big') +
                    int(hex_tail, 16).to_bytes(8, 'big'))
        self.send_data(data_bin)

    def send_FESConfig_to_queue(self, fes_Config):
        """
        将FES配置信息发送到下位机
        """
        logger.info("------FES-----")
        data1 = 0x55
        # data1 = int(hex_num, 16)
        hex_tail = '0x11'
        data_bin = (data1.to_bytes(4, 'big') +
                    fes_Config.STI_ch.to_bytes(2, 'big') +
                    fes_Config.ele_Intensity.to_bytes(2, 'big') +
                    fes_Config.Worktime.to_bytes(2, 'big') +
                    fes_Config.Lastingtime.to_bytes(1, 'big') +
                    fes_Config.HoldOntime.to_bytes(1, 'big') +
                    fes_Config.Waittingtime.to_bytes(2, 'big') +
                    fes_Config.WaveNum.to_bytes(2, 'big')
                    )
        self.send_data(data_bin)

    def send_fescongfig_intensity(self, intensity, ch=0):
        tmp_fes_config = FES_config_map[intensity]
        tmp_fes_config.STI_ch = ch
        self.send_FESConfig_to_queue(tmp_fes_config)

    # todo : 设备端开启FES功能
    def beginFES(self):
        hex_num = '0x66'
        self.send_data(int('0x66', 16).to_bytes(4, 'big'))

    def Test_send_FESConfig(self):
        fes_Config1 = FES_Config(ch=0, Worktime=2000, Lastingtime=150, HoldOntime=48)
        self.send_FESConfig_to_queue(fes_Config1)
        fes_Config1.set_FESWave_Form(500, 2)
        self.send_FESConfig_to_queue(fes_Config1)

    def open_ele(self):
        self.send_data(int('0x66', 16).to_bytes(4, 'big') +
                       int('0x1141', 16).to_bytes(12, 'big'))

    def switch_debug_mode(self):
        self.send_data(int('0x77', 16).to_bytes(4, 'big') +
                       int('0x1141', 16).to_bytes(12, 'big'))


class MainWindow(QMainWindow):
    def __init__(self):
        super(MainWindow, self).__init__()
        # 额外开启的进程
        self.stop_event = None
        self.process_Hand = None
        self.record_process = None

        self.ConnectionFlag = 0
        self.ui = Ui_MainWindow()
        self.ui.setupUi(self)  # 设置UI

        # 添加图形对象
        # 创建一个网格布局
        layout = QGridLayout()
        self.ui.graph_widget.setLayout(layout)
        self.plots = [pg.PlotWidget() for _ in range(4)]
        self.curves = []

        for i, plot in enumerate(self.plots):
            # plot.plot(x, y, pen='r')
            self.curves.append(plot.plot())
            plot.setBackground('w')
            plot.setLabel('left', 'sEMG', 'mV')
            plot.setLabel('bottom', 'Time', 'ms')
            plot.setYRange(0, 3500)
            # title
            plot.setTitle("sEMG ch" + str(i))

        # 将图形添加到布局中
        layout.addWidget(self.plots[0], 0, 0)
        layout.addWidget(self.plots[1], 0, 1)
        layout.addWidget(self.plots[2], 1, 0)
        layout.addWidget(self.plots[3], 1, 1)

        # 这里可以添加更多的初始化代码和信号槽连接
        # todo ： 添加信号槽连接
        self.ui.pushButton_beginconnect.pressed.connect(lambda: self.begin_connect(self.ui.pushButton_beginconnect))
        self.ui.pushButton_record_sEMG.pressed.connect(
            lambda: self.record_data(self.ui.pushButton_record_sEMG.isChecked()))
        self.ui.pushButton_send.pressed.connect(self.Send_command)
        self.ui.comboBox_sEMG.currentIndexChanged.connect(
            lambda: self.Set_sEMG_mode(self.ui.comboBox_sEMG.currentIndex()))
        self.ui.comboBox_controlmode.currentIndexChanged.connect(
            lambda: self.Control_mode_changed(self.ui.comboBox_controlmode.currentIndex()))
        #UDP地址配置
        self.ui.pushButton_beginconnect.pressed.connect(lambda: self.Set_UDP_addr())

        # 绑定sEMGch复选框勾选事件
        self.ui.checkBox_ch0.stateChanged.connect(lambda: self.Set_ch_state(0, self.ui.checkBox_ch0.isChecked()))
        self.ui.checkBox_ch1.stateChanged.connect(lambda: self.Set_ch_state(1, self.ui.checkBox_ch1.isChecked()))
        self.ui.checkBox_ch2.stateChanged.connect(lambda: self.Set_ch_state(2, self.ui.checkBox_ch2.isChecked()))
        self.ui.checkBox_ch3.stateChanged.connect(lambda: self.Set_ch_state(3, self.ui.checkBox_ch3.isChecked()))

        self.ui.pushButton_sendFESconfig.pressed.connect(self.Send_FES_config)
        self.ui.CH_Slider_1.valueChanged.connect(lambda: self.CH_Slider_changed(0, self.ui.CH_Slider_1))
        self.ui.CH_Slider_2.valueChanged.connect(lambda: self.CH_Slider_changed(1, self.ui.CH_Slider_2))
        self.ui.CH_Slider_3.valueChanged.connect(lambda: self.CH_Slider_changed(2, self.ui.CH_Slider_3))
        self.ui.CH_Slider_4.valueChanged.connect(lambda: self.CH_Slider_changed(3, self.ui.CH_Slider_4))

        # 初始化UDP通信和数据处理对象
        self.udp_communicator = UDPCommunicator(LocalHost, RemoteHost, 8080)

        # 初始化FES控制对象，传递 communicator 实例和 RMS 输入队列
        self.fes_controler = FES_Controler(self.udp_communicator.rms_for_fes_queue,
                                           self.udp_communicator.FES_config_queue)
        self.fes_control_thread = threading.Thread(target=self.fes_controler.control_thread_loop)

        # 启动数据接收和处理线程
        self.data_thread = threading.Thread(target=self.udp_communicator.data_process_loop)
        self.udp_thread = threading.Thread(target=self.udp_communicator.receive_and_send_loop)

        self.data_thread.start()
        self.udp_thread.start()

        # 创建定时器，用于更新图表
        self.timer = QtCore.QTimer()
        self.timer.timeout.connect(self.update_plots)
        self.timer.start(200)  # 每隔200ms更新一次图表

        self.openfile = None

        self.timer2 = QtCore.QTimer()
        self.timer2.timeout.connect(self.udp_communicator.send_FES_control)
        self.timer2.start(FES_FREQ)  # 电刺激指令每隔FES_FREQms执行一次

    # 初始化完毕

    # todo: 槽函数
    # 槽函数定义
    def begin_connect(self, button):
        if not button.isChecked():
            self.ConnectionFlag = 1
            self.shell_display("开始连接")
        else:
            self.ConnectionFlag = 0
            self.shell_display("断开连接")

    def record_data(self, isChecked):
        # done:多进程写文件
        global Data_Writeflag
        if not isChecked:
            self.shell_display("开始录制")
            Data_Writeflag = True
            # 开启录制进程
            self.stop_event = multiprocessing.Event()
            self.record_process = multiprocessing.Process(target=data_recorder.sEMG_record_process,
                                                          args=(WriteIn_Queue, self.stop_event))
            if sensors.six_axis_force:
                self.process_Hand = multiprocessing.Process(target=mainProcess_Task_six_axis_force)
                self.process_Hand.start()
            self.record_process.start()
            self.shell_display("开启录制进程")
        else:
            Data_Writeflag = False
            # 发送关闭请求
            # 通过Event通知停止
            self.stop_event.set()
            if sensors.six_axis_force:
                self.process_Hand.terminate()
            # 等待进程正常退出
            self.record_process.join(timeout=1)
            self.shell_display("停止录制")

    def Send_command(self):
        command = self.ui.lineEdit_command.text()
        self.shell_display("发送指令: " + command)
        if self.udp_communicator.addr is not None:
            self.udp_communicator.CommandExecute(command)
        else:
            self.shell_display("未连接")

    def Set_sEMG_mode(self, curindex):
        global sEMG_mode
        if curindex == 0:
            sEMG_mode = 0
            self.shell_display("设置sEMG模式为：raw sEMG")
        elif curindex == 1:
            sEMG_mode = 1
            self.shell_display("设置sEMG模式为：RMS")

    def shell_display(self, show_string):
        self.ui.textBrowser.append(show_string)

    def Set_ch_state(self, ch, Ischecked):
        global channel_state
        channel_state.sEMGChannel_State[ch] = Ischecked
        self.shell_display("通道" + str(ch) + "状态：" + str(Ischecked))

    def Control_mode_changed(self, curindex):
        self.shell_display("控制模式切换为：" + str(curindex))
        self.fes_controler.control_mode = curindex
        global FES_Control_ON, FES_FREQ
        if curindex == 0:
            # 关闭FES自动控制
            FES_Control_ON = False
            # 关闭QTimer2
            self.timer2.stop()
            # 发送停止信号
            self.fes_controler.stop_event.set()
        else:
            FES_Control_ON = True

            self.timer2.start(FES_FREQ)
            # 发送启动信号
            self.fes_controler.stop_event.clear()
            if not self.fes_control_thread.is_alive():
                self.fes_control_thread.start()

            self.shell_display("开启FES控制线程")

    def CH_Slider_changed(self, ch, slider):
        channel_state.channel_intensity[ch] = slider.value()
        channel_state.FESchanged_state[ch] = True

    def Send_FES_config(self):
        for i in range(4):
            if channel_state.FESchanged_state[i]:
                channel_state.FESchanged_state[i] = False
                self.udp_communicator.SetTensity(i, channel_state.channel_intensity[i])

    def update_plots(self):
        # data = self.data_processor.get_data_for_plot()
        """
        测试图像绘制
        控制窗口最大数据量
        """
        # data = [np.random.randn(100) for _ in range(4)]
        for (i, curve) in enumerate(self.curves):
            # curve.setData(data[i])  #测试模式
            if channel_state.sEMGChannel_State[i]:
                curve.setData(self.udp_communicator.plot_data_list[i])
        # global FES_Control_ON, FES_State
        # if FES_Control_ON:
        #     curve.setPen('r')
        #     for (i, curve) in enumerate(self.curves):
        #         # 绘制电刺激曲线,红色
        #         if FES_State[i]:
        #             # 画一段电刺激双极性曲线
        #             curve.setData(FES_State.FES_line)
        #     curve.setPen('b')


def main():
    app = QApplication(sys.argv)
    mainWin = MainWindow()
    qtmodern.styles.dark(app)
    mw = qtmodern.windows.ModernWindow(mainWin)
    mw.show()
    sys.exit(app.exec())


if __name__ == "__main__":
    app = QApplication(sys.argv)
    mainWin = MainWindow()
    qtmodern.styles.dark(app)
    mw = qtmodern.windows.ModernWindow(mainWin)
    mw.show()
    sys.exit(app.exec())
